addpath ../
addpath ../export_fig/


% ****************************************************************
% test imputation methods on RAW log expression
% 
ntis = 53;
ngene = 1000;
nind = 450;
nfactor = 15;
nconfounder = 5;
factor_per_gene = 3;
pmiss = .8;

U_true = randn(ntis,nfactor,'single');
Y_true = zeros(ntis,nind,ngene,'single');

for g = 1:ngene,
    V = randn(factor_per_gene,nind,'single');
    yy = zscore(U_true(:,randsample(nfactor,factor_per_gene))*V,[],2);
    Y_true(:,:,g) = yy;
    clear yy;
end

% Add confounder effects
% & introduce missing individuals per tissue
C_true = randn(nconfounder, nind, 'single');
Y_train = NaN(ntis,nind,ngene,'single');

for t = 1:ntis,
    Y = squeeze(Y_true(t,:,:)); % ind x gene
    W = randn(ngene,nconfounder,'single');
    Y = zscore(Y + (W*C_true)');    

    missing_ind = randsample(nind, floor(nind * pmiss));
    Y(missing_ind,:) = NaN;

    Y_train(t,:,:) = Y;
end




% ****************************************************************
% confounder correction
Y_train_corrected = Y_train;
J = 10;
for t = 1:ntis,
    Y = squeeze(Y_train_corrected(t,:,:));
    obs_ind = sum(isnan(Y')) == 0;
    [Uavg,Vavg,tau,rmse] = run_modQTL_confounder(Y(obs_ind,:)',J,10,1000);

    Yc = Y(obs_ind,:)' - Uavg*Vavg;
    Y_train_corrected(t,obs_ind,:) = Yc';
end



% ****************************************************************
J = 30;
burnin = 100;
nepoch = 900;
batchsize = 500;


outdir = 'sim02_result';
run_modQTL_imputation_svi_vect(Y_train,J,burnin,nepoch,batchsize,outdir);

outdir_corrected = 'sim02_result_corrected';
run_modQTL_imputation_svi_vect(Y_train_corrected,J,burnin,nepoch,batchsize,outdir_corrected);


% ****************************************************************
load([outdir,'/parameters.mat']);

files = dir([outdir,'/V_minibatch*.mat']);
yytrue = [];
yyhat = [];

for ff = 1:numel(files),
    load([outdir,'/',files(ff).name]);
    
    Y = Y_true(:,:,genes_minibatch);
    missing = isnan(Y_train(:,:,genes_minibatch));

    yytrue = [yytrue; Y(missing)];

    YY = zeros(ntis,nind,numel(genes_minibatch));

    for g = 1:numel(genes_minibatch),
        YY(:,:,g) = U * V_minibatch(:,:,g);
    end

    yyhat = [yyhat; YY(missing)];

    clear Y *_minibatch;
end

[r,p] = corr(yytrue,yyhat);

close all;
h1 = figure(1);
set(h1,'position',[.5 .5 400 300]);
idx = randsample(numel(yytrue),10000);
hold on;
dscatter( yytrue(idx), yyhat(idx) );
hold off;
xlabel('true'); ylabel('imputed');
title(sprintf('r = %.2e p = %.2e',r,p));
axis('tight');

export_fig('sim02_fig01','-pdf','-a1','-m2','-transparent');


% ****************************************************************
load([outdir_corrected,'/parameters.mat']);

files = dir([outdir_corrected,'/V_minibatch*.mat']);
yytrue = [];
yyhat = [];

for ff = 1:numel(files),
    load([outdir_corrected,'/',files(ff).name]);
    
    Y = Y_true(:,:,genes_minibatch);
    missing = isnan(Y_train(:,:,genes_minibatch));

    yytrue = [yytrue; Y(missing)];

    YY = zeros(ntis,nind,numel(genes_minibatch));

    for g = 1:numel(genes_minibatch),
        YY(:,:,g) = U * V_minibatch(:,:,g);
    end

    yyhat = [yyhat; YY(missing)];

    clear Y *_minibatch;
end

[r,p] = corr(yytrue,yyhat);

close all;
h1 = figure(1);
set(h1,'position',[.5 .5 400 300]);
idx = randsample(numel(yytrue),10000);
hold on;
dscatter( yytrue(idx), yyhat(idx) );
plot( [-3,3], r*[-3,3], 'r--', 'linewidth', 2 );
plot( [-3,3], [-3,3], 'k--', 'linewidth', 2 );
hold off;
xlabel('true'); ylabel('imputed');
title(sprintf('r = %.2e p = %.2e',r,p));
axis('tight');

export_fig('sim02_fig02','-pdf','-a1','-m2','-transparent');


